from .resources import Directory, Account, Order, Authorization, Challenge
import requests
from cryptography import x509
from cryptography.hazmat.primitives.asymmetric import rsa, ec, types
from cryptography.hazmat.primitives import hashes, serialization
import hashlib
from jwcrypto import jwk, jws
from jwcrypto.common import json_encode, json_decode, base64url_encode
import pathlib
from typing import Any, List, Union

_CSRHashTypes = Union[
    hashes.SHA224,
    hashes.SHA256,
    hashes.SHA384,
    hashes.SHA512,
    hashes.SHA3_224,
    hashes.SHA3_256,
    hashes.SHA3_384,
    hashes.SHA3_512,
]

class SimpleAcmeClient:
    USER_AGENT = "simpleacme-python"
    JSON_HEADER = "application/json"
    JOSE_HEADER = "application/jose+json"

    def __init__(
        self,
        directory: Directory,
        key: jwk.JWK | pathlib.Path,
        key_password: str | bytes = None,
        signature_alg: str = "RS256",
    ):
        """构建一个SimpleAcme客户端

        Args:
            directory (Directory): ACME服务器的Directory，可以使用类方法SimpleAcmeClient.get_directory获取
            key (jwk.JWK | pathlib.Path): jwcrypto的JWK类或者PEM文件的路径
            key_password (str | bytes, optional): 若要使用的密钥设有密码，请指定此参数
            signature_alg (str, optional): 默认使用RS256进行签名，若要使用其他签名算法，请指定此参数，可用签名方式参见 https://www.iana.org/assignments/jose/jose.xhtml#web-signature-encryption-algorithms中的可用alg名称

        Raises:
            ValueError: PEM文件有误
            Exception: 密钥不完整（缺少私钥或公钥）
        """
        self.directory = directory
        self.nonce = set()
        self.signature_alg = signature_alg

        # 导入密钥
        if isinstance(key, jwk.JWK):
            self._key = key
        elif isinstance(key, pathlib.Path):
            if isinstance(key_password, str):
                psw = key_password.encode("utf-8")
            else:
                psw = key_password
            with open(key, "rb") as pem_file:
                try:
                    self._key = jwk.JWK()
                    self._key.import_from_pem(pem_file.read(), psw)
                except ValueError:
                    raise ValueError("PEM file incorrect, import failed")
        # 检查密钥是否完整
        if not self._key.has_private:
            raise Exception("No private key provided")
        if not self._key.has_public:
            raise Exception("No public key provided")

    def download_certificate(self, order: Order) -> bytes:
        protected = {
            "alg": self.signature_alg,
            "kid": self.account.account_url,
            "nonce": self._get_nonce(),
            "url": order.certificate,
        }

        # 添加protected并完成签名，此处为POST-as-GET操作，payload为空串即可
        jws_obj = jws.JWS("".encode("utf-8"))
        jws_obj.add_signature(self._key, alg=self.signature_alg, protected=protected)
        try:
            resp = requests.post(
                order.certificate,
                json=json_decode(jws_obj.serialize()),
                headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
            )
            # 添加Nonce
            if "Replay-Nonce" in resp.headers:
                self.nonce.add(resp.headers["Replay-Nonce"])
        except requests.exceptions.RequestException as e:
            raise Exception(str(e))

        if resp.status_code == 200:
            return resp.content
        else:
            raise Exception(f"{resp.status_code}: {resp.text}")

    def start_finalization(
        self,
        csr_builder: x509.CertificateSigningRequestBuilder,
        order: Order,
        signing_key: types.CertificateIssuerPrivateKeyTypes,
        signing_algo: _CSRHashTypes
    ):
        # signing_key私钥用于对CSR（Certificate Signing Request）进行签名
        csr = csr_builder.sign(signing_key, signing_algo())
        csr_der = csr.public_bytes(serialization.Encoding.DER)

        protected = {
            "alg": self.signature_alg,
            "kid": self.account.account_url,
            "nonce": self._get_nonce(),
            "url": order.finalize,
        }
        payload = {"csr": base64url_encode(csr_der)}

        jws_obj = jws.JWS(json_encode(payload).encode("utf-8"))
        jws_obj.add_signature(self._key, alg=self.signature_alg, protected=protected)
        try:
            resp = requests.post(
                order.finalize,
                json=json_decode(jws_obj.serialize()),
                headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
            )

            # 响应体bytes -> 按utf8解码的str -> Python object
            resp_decoded = json_decode(resp.content.decode("utf-8"))
            # 添加Nonce
            if "Replay-Nonce" in resp.headers:
                self.nonce.add(resp.headers["Replay-Nonce"])
        except requests.exceptions.RequestException as e:
            raise Exception(str(e))

        if resp.status_code == 200:
            return Order(resp_decoded, resp.headers["Location"])
        else:
            raise Exception(f"{resp.status_code}: {resp.text}")

    def answer_challenge(self, authorization: Authorization, chall_type: str):
        chall = None
        for challenge in authorization.challenges:
            if challenge["type"] == chall_type:
                chall = challenge
                break
        if chall is None:
            raise ValueError("challenge type is invalid")

        protected = {
            "alg": self.signature_alg,
            "kid": self.account.account_url,
            "nonce": self._get_nonce(),
            "url": chall["url"],
        }

        # 此处的payload是空的json，即解析为字符串后为"{}"
        jws_obj = jws.JWS(json_encode({}).encode("utf-8"))
        jws_obj.add_signature(self._key, alg=self.signature_alg, protected=protected)
        try:
            resp = requests.post(
                chall["url"],
                json=json_decode(jws_obj.serialize()),
                headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
            )

            # 响应体bytes -> 按utf8解码的str -> Python object
            resp_decoded = json_decode(resp.content.decode("utf-8"))
            # 添加Nonce
            if "Replay-Nonce" in resp.headers:
                self.nonce.add(resp.headers["Replay-Nonce"])
        except requests.exceptions.RequestException as e:
            raise Exception(str(e))

        if resp.status_code == 200:
            return Challenge(resp_decoded)
        else:
            raise Exception(f"{resp.status_code}: {resp.text}")

    def new_authorization(self, order: Order) -> List[Authorization]:
        res = []
        for authz in order.authorizations:
            # 构建JWT的protected部分
            protected = {
                "alg": self.signature_alg,
                "kid": self.account.account_url,
                "nonce": self._get_nonce(),
                "url": authz,
            }
            # 添加protected并完成签名，此处为POST-as-GET操作，payload为空串即可
            jws_obj = jws.JWS("".encode("utf-8"))
            jws_obj.add_signature(
                self._key, alg=self.signature_alg, protected=protected
            )
            try:
                resp = requests.post(
                    authz,
                    json=json_decode(jws_obj.serialize()),
                    headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
                )

                # 响应体bytes -> 按utf8解码的str -> Python object
                resp_decoded = json_decode(resp.content.decode("utf-8"))
                # 添加Nonce
                if "Replay-Nonce" in resp.headers:
                    self.nonce.add(resp.headers["Replay-Nonce"])
            except requests.exceptions.RequestException as e:
                raise Exception(str(e))

            if resp.status_code == 200 or resp.status_code == 201:
                res.append(Authorization(resp_decoded))
            else:
                raise Exception(f"{resp.status_code}: {resp.text}")

        return res

    def convert_token(self, token: str) -> str:
        challenge_str = (token + "." + self._key.thumbprint()).encode("utf-8")
        challenge_raw = hashlib.sha256(challenge_str).digest()
        challenge_final = base64url_encode(challenge_raw)
        return challenge_final

    def new_order(self, domains: List[str]):
        account = getattr(self, "account", None)
        if account is None:
            raise Exception("No account registered, call get_account first")
        # 构建JWT的payload部分
        payload = {
            "identifiers": [{"type": "dns", "value": domain} for domain in domains]
        }
        # 构建JWT的protected部分
        protected = {
            "alg": self.signature_alg,
            "kid": account.account_url,
            "nonce": self._get_nonce(),
            "url": self.directory.newOrder,
        }
        # 添加protected并完成签名
        jws_obj = jws.JWS(json_encode(payload).encode("utf-8"))
        jws_obj.add_signature(self._key, alg=self.signature_alg, protected=protected)
        try:
            resp = requests.post(
                self.directory.newOrder,
                json=json_decode(jws_obj.serialize()),
                headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
            )
            # 响应体bytes -> 按utf8解码的str -> Python object
            resp_decoded = json_decode(resp.content.decode("utf-8"))
            # 添加Nonce
            if "Replay-Nonce" in resp.headers:
                self.nonce.add(resp.headers["Replay-Nonce"])
        except requests.exceptions.RequestException as e:
            raise Exception(str(e))

        if resp.status_code == 201:
            order = Order(resp_decoded, resp.headers["Location"])
            return order
        else:
            raise Exception(f"{resp.status_code}: {resp.text}")

    def get_account(
        self,
        termsOfServiceAgreed: bool,
        onlyReturnExisting: bool = False,
        externalAccountBinding: Any = None,
        contact: List[str] = None,
    ):
        if not termsOfServiceAgreed:
            return

        # 构建JWT的payload部分
        payload = {"termsOfServiceAgreed": termsOfServiceAgreed}
        if contact:
            payload["contact"] = contact
        if onlyReturnExisting:
            payload["onlyReturnExisting"] = onlyReturnExisting
        if externalAccountBinding:
            payload["externalAccountBinding"] = externalAccountBinding
        # 构建JWT的protected部分
        protected = {
            "alg": self.signature_alg,
            "jwk": self._key.export_public(as_dict=True),
            "nonce": self._get_nonce(),
            "url": self.directory.newAccount,
        }
        # 添加protected并完成签名
        jws_obj = jws.JWS(json_encode(payload).encode("utf-8"))
        jws_obj.add_signature(self._key, alg=self.signature_alg, protected=protected)

        # 发起请求
        try:
            # 此处先将序列化后的整个请求体decode以便作为requests的json参数
            resp = requests.post(
                self.directory.newAccount,
                json=json_decode(jws_obj.serialize()),
                headers={"Content-Type": SimpleAcmeClient.JOSE_HEADER},
            )
            # 响应体bytes -> 按utf8解码的str -> Python object
            resp_decoded = json_decode(resp.content.decode("utf-8"))
            # 添加Nonce
            if "Replay-Nonce" in resp.headers:
                self.nonce.add(resp.headers["Replay-Nonce"])
        except requests.exceptions.RequestException as e:
            raise Exception(str(e))

        # 201为新创建用户、200为查找到已存在用户
        if resp.status_code in (200, 201):
            resp_decoded["termsOfServiceAgreed"] = True
            self.account = Account(resp_decoded, resp.headers["Location"])
            return self.account
        else:
            raise Exception(f"{resp.status_code}: {resp.text}")

    def _get_nonce(self) -> str:
        """获取nonce

        没有库存nonce时进行请求

        Returns:
            str: 可用的nonce
        """
        if self.nonce:
            return self.nonce.pop()
        else:
            try:
                resp = requests.head(self.directory.newNonce)
            except requests.exceptions.RequestException as e:
                raise Exception(str(e))
            if resp.status_code != 200:
                raise Exception(f"{resp.status_code}: {resp.text}")
            return resp.headers["Replay-Nonce"]
    
    def refresh_resource(self, resource: Order | Authorization):
        pass

    @classmethod
    def get_directory(cls, directory_url: str) -> Directory:
        """根据ACME服务器的Directory地址获取整个Directory

        Args:
            directory_url (str): ACME服务器的Directory地址

        Raises:
            Exception: 网络请求错误时抛出

        Returns:
            Directory: Directory的示例，包含了服务器返回的各资源地址
        """
        try:
            resp = requests.get(
                directory_url,
                headers={"Content-Type": cls.JSON_HEADER, "User-Agent": cls.USER_AGENT},
            )
            # 响应体bytes -> 按utf8解码的str -> Python object
            resp_decoded = json_decode(resp.content.decode("utf-8"))
        except requests.exceptions.RequestException as e:
            raise
        if resp.status_code != 200:
            raise Exception(f"{resp.status_code}: {resp.text}")

        return Directory(resp_decoded)
